package com.zenithmind.chat.service.impl;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.knuddels.jtokkit.Encodings;
import com.knuddels.jtokkit.api.Encoding;
import com.knuddels.jtokkit.api.EncodingRegistry;
import com.knuddels.jtokkit.api.EncodingType;
import com.zenithmind.chat.config.AiConfig;
import com.zenithmind.chat.pojo.entity.AiChat;
import com.zenithmind.chat.pojo.entity.AiChatMessage;
import com.zenithmind.chat.pojo.entity.ModelConfig;
import com.zenithmind.chat.service.AiService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;


@Slf4j
@Service
@Qualifier("ollamaService")
public class OllamaServiceImpl implements AiService {
    
    @Autowired
    private AiConfig aiConfig;
    
    @Autowired
    private RestTemplate restTemplate;
    
    @Autowired
    private ObjectMapper objectMapper;
    
    // 创建编码注册表实例（单例）
    private static final EncodingRegistry registry = Encodings.newDefaultEncodingRegistry();
    
    @Override
    public AiChatMessage chat(AiChat chat, List<AiChatMessage> messages, String content) {
        try {
            // 获取模型配置
            ModelConfig modelConfig = aiConfig.getModelConfig(2, chat.getModelName());
            
            // 如果数据库中没有配置，则使用默认配置
            String apiUrl = aiConfig.getOllama().getApiUrl();
            
            // 如果数据库中有配置，则使用数据库配置
            if (modelConfig != null && modelConfig.getApiUrl() != null && !modelConfig.getApiUrl().isEmpty()) {
                apiUrl = modelConfig.getApiUrl();
            }
            
            // 构建请求头
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            
            // 构建消息列表
            List<Map<String, String>> messageList = new ArrayList<>();
            
            // 添加系统提示词
            if (chat.getSystemPrompt() != null && !chat.getSystemPrompt().isEmpty()) {
                Map<String, String> systemMessage = new HashMap<>();
                systemMessage.put("role", "system");
                systemMessage.put("content", chat.getSystemPrompt());
                messageList.add(systemMessage);
            }
            
            // 添加历史消息
            for (AiChatMessage message : messages) {
                Map<String, String> messageMap = new HashMap<>();
                messageMap.put("role", message.getRole());
                messageMap.put("content", message.getContent());
                messageList.add(messageMap);
            }
            
            // 添加用户新消息
            Map<String, String> userMessage = new HashMap<>();
            userMessage.put("role", "user");
            userMessage.put("content", content);
            messageList.add(userMessage);
            
            // 构建请求体
            Map<String, Object> requestBody = new HashMap<>();
            requestBody.put("model", chat.getModelName());
            requestBody.put("messages", messageList);
            requestBody.put("stream", false);
            
            // 发送请求
            String response = restTemplate.postForObject(
                apiUrl + "/api/chat",
                new HttpEntity<>(requestBody, headers),
                String.class
            );
            
            // 解析响应
            Map<String, Object> responseMap = objectMapper.readValue(response, Map.class);
            Map<String, String> message = (Map<String, String>) responseMap.get("message");
            
            // 获取token计数（如果存在）
            int tokenCount;
            if (responseMap.containsKey("total_duration") && responseMap.containsKey("eval_count")) {
                // 有些Ollama模型会返回eval_count，这与token数相关
                tokenCount = ((Number)responseMap.get("eval_count")).intValue();
            } else {
                // 否则估算
                tokenCount = countTokens(message.get("content"));
            }
            
            // 创建AI回复消息
            AiChatMessage aiMessage = new AiChatMessage();
            aiMessage.setChatId(chat.getId());
            aiMessage.setRole("assistant");
            aiMessage.setContent(message.get("content"));
            aiMessage.setTokenCount(tokenCount);
            aiMessage.setStatus(1);
            
            return aiMessage;
            
        } catch (Exception e) {
            log.error("Ollama API调用失败", e);
            throw new RuntimeException("AI服务暂时不可用，请稍后重试");
        }
    }
    
    @Override
    public AiChatMessage chatStream(AiChat chat, List<AiChatMessage> messages, String content, SseEmitter emitter) {
        // 获取模型配置
        ModelConfig modelConfig = aiConfig.getModelConfig(2, chat.getModelName());
        String apiUrl = (modelConfig != null && modelConfig.getApiUrl() != null && !modelConfig.getApiUrl().isEmpty()) 
                        ? modelConfig.getApiUrl() 
                        : aiConfig.getOllama().getApiUrl();

        // 构建请求头
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        // Ollama stream API typically returns application/x-ndjson or similar line-delimited JSON
        // headers.setAccept(List.of(MediaType.APPLICATION_NDJSON, MediaType.APPLICATION_JSON)); 
        // Simpler to let RestTemplate handle it with default converters for InputStream

        List<Map<String, String>> messageList = new ArrayList<>();
        if (chat.getSystemPrompt() != null && !chat.getSystemPrompt().isEmpty()) {
            messageList.add(Map.of("role", "system", "content", chat.getSystemPrompt()));
        }
        for (AiChatMessage message : messages) {
            messageList.add(Map.of("role", message.getRole(), "content", message.getContent()));
        }
        messageList.add(Map.of("role", "user", "content", content));

        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", chat.getModelName());
        requestBody.put("messages", messageList);
        requestBody.put("stream", true);

        StringBuilder fullResponseContent = new StringBuilder();
        final int[] totalEvalCount = {0};
        
        BiConsumer<String, String> sendSseEvent = (type, eventContent) -> {
            try {
                Map<String, String> eventMap = new HashMap<>();
                eventMap.put("type", type);
                eventMap.put("content", eventContent);
                
                // 思考内容使用INFO级别记录，确保在控制台可见
                if ("thought".equals(type)) {
                    log.info("Sending thought event: content={}", eventContent.substring(0, Math.min(50, eventContent.length())) + (eventContent.length() > 50 ? "..." : ""));
                } else {
                    log.debug("Sending SSE event: type={}, content={}", type, eventContent.substring(0, Math.min(50, eventContent.length())) + (eventContent.length() > 50 ? "..." : ""));
                }
                
                emitter.send(SseEmitter.event().data(objectMapper.writeValueAsString(eventMap)));
            } catch (Exception e) {
                log.warn("Failed to send SSE event: type={}, content={}", type, eventContent, e);
            }
        };

        try {
            // Simulate initial "thought" process
            sendSseEvent.accept("thought", "正在连接到 " + chat.getModelName() + " 模型...");
            // Simulate another "thought"
            sendSseEvent.accept("thought", "模型正在分析您的请求...");
            
            // 设置超时时间更短的RestTemplate，以便更快捕获超时问题
            RestTemplate timeoutRestTemplate = new RestTemplate();
            // 设置连接超时时间为10秒，读取超时为30秒
            timeoutRestTemplate.setRequestFactory(new org.springframework.http.client.SimpleClientHttpRequestFactory() {
                @Override
                public void setConnectTimeout(int timeout) {
                    super.setConnectTimeout((int) java.time.Duration.ofSeconds(10).toMillis());
                }
                
                @Override
                public void setReadTimeout(int timeout) {
                    super.setReadTimeout((int) java.time.Duration.ofSeconds(30).toMillis());
                }
            });

            try {
                timeoutRestTemplate.execute(
                    apiUrl + "/api/chat",
                    org.springframework.http.HttpMethod.POST,
                    request -> {
                        request.getHeaders().putAll(headers);
                        // Ensure RestTemplate sends request body as UTF-8
                        request.getHeaders().setContentType(MediaType.APPLICATION_JSON_UTF8); 
                        objectMapper.writeValue(request.getBody(), requestBody);
                    },
                    response -> {
                        try (java.io.InputStream inputStream = response.getBody();
                             java.io.BufferedReader reader = new java.io.BufferedReader(new java.io.InputStreamReader(inputStream, java.nio.charset.StandardCharsets.UTF_8))) {
                            String line;
                            StringBuilder inProgressContent = new StringBuilder();
                            boolean inThinkSection = false;
                            StringBuilder thinkContent = new StringBuilder();
                            
                            while ((line = reader.readLine()) != null) {
                                try {
                                    Map<String, Object> chunkMap = objectMapper.readValue(line, Map.class);
                                    if (chunkMap.containsKey("message")) {
                                        Map<String, String> messageChunkMap = (Map<String, String>) chunkMap.get("message");
                                        String contentChunk = messageChunkMap.get("content");
                                        
                                        if (contentChunk != null) {
                                            // Always append the raw chunk to fullResponseContent first
                                            fullResponseContent.append(contentChunk);

                                            inProgressContent.append(contentChunk);
                                            String currentContent = inProgressContent.toString();
                                            
                                            // 处理<think>标签
                                            if (currentContent.contains("<think>") && !inThinkSection) {
                                                // 切换到思考模式
                                                inThinkSection = true;
                                                // 提取<think>标签前的内容(如果有)作为回答
                                                int thinkStart = currentContent.indexOf("<think>");
                                                if (thinkStart > 0) {
                                                    String answerBeforeThink = currentContent.substring(0, thinkStart);
                                                    sendSseEvent.accept("answer_chunk", answerBeforeThink);
                                                    // fullResponseContent.append(answerBeforeThink); // Already appended
                                                }
                                                thinkContent = new StringBuilder();
                                                // 提取思考内容(如果<think>后面已经有内容)
                                                if (currentContent.length() > thinkStart + 7) { // "<think>".length() = 7
                                                    String initialThinkContent = currentContent.substring(thinkStart + 7);
                                                    thinkContent.append(initialThinkContent);
                                                    
                                                    // 发送第一个思考块，实现真正的流式思考输出
                                                    sendSseEvent.accept("thought", initialThinkContent);
                                                    
                                                    if (initialThinkContent.contains("</think>")) {
                                                        // 如果同一个chunk中包含完整的<think>和</think>
                                                        int thinkEnd = initialThinkContent.indexOf("</think>");
                                                        
                                                        // 切换回回答模式
                                                        inThinkSection = false;
                                                        
                                                        // 如果</think>后面还有内容，发送为回答
                                                        if (initialThinkContent.length() > thinkEnd + 8) { // "</think>".length() = 8
                                                            String answerAfterThink = initialThinkContent.substring(thinkEnd + 8);
                                                            sendSseEvent.accept("answer_chunk", answerAfterThink);
                                                            // fullResponseContent.append(answerAfterThink); // Already appended
                                                        }
                                                        
                                                        // 重置累积内容，因为已经处理完一对<think></think>标签
                                                        inProgressContent = new StringBuilder(
                                                            initialThinkContent.substring(Math.min(initialThinkContent.length(), thinkEnd + 8))
                                                        );
                                                    }
                                                }
                                            } 
                                            else if (inThinkSection) {
                                                // 在思考模式中，累积思考内容
                                                thinkContent.append(contentChunk);
                                                
                                                // 将每个思考块直接发送，实现真正的流式思考输出
                                                sendSseEvent.accept("thought", contentChunk);
                                                
                                                String currentThinkContent = thinkContent.toString();
                                                
                                                // 检查是否有</think>标签
                                                if (currentThinkContent.contains("</think>")) {
                                                    int thinkEnd = currentThinkContent.indexOf("</think>");
                                                    
                                                    // 切换回回答模式
                                                    inThinkSection = false;
                                                    
                                                    // 如果</think>后面还有内容，发送为回答
                                                    if (currentThinkContent.length() > thinkEnd + 8) {
                                                        String answerAfterThink = currentThinkContent.substring(thinkEnd + 8);
                                                        sendSseEvent.accept("answer_chunk", answerAfterThink);
                                                        // fullResponseContent.append(answerAfterThink); // Already appended
                                                    }
                                                    
                                                    // 重置累积内容，只保留</think>后尚未处理的部分
                                                    inProgressContent = new StringBuilder(
                                                        currentThinkContent.substring(Math.min(currentThinkContent.length(), thinkEnd + 8))
                                                    );
                                                }
                                            } 
                                            else {
                                                // 普通回答模式，发送内容块
                                                sendSseEvent.accept("answer_chunk", contentChunk);
                                                // fullResponseContent.append(contentChunk); // Already appended
                                            }
                                        }
                                    }
                                    if (chunkMap.containsKey("done") && (Boolean) chunkMap.get("done")) {
                                        if (chunkMap.containsKey("eval_count")) {
                                            totalEvalCount[0] = ((Number)chunkMap.get("eval_count")).intValue();
                                        }
                                        
                                        // 如果结束时仍在思考模式，发送剩余的思考内容
                                        if (inThinkSection && thinkContent.length() > 0) {
                                            sendSseEvent.accept("thought", thinkContent.toString());
                                        }
                                        
                                        sendSseEvent.accept("thought", "模型已完成响应生成。");
                                        break; 
                                    }
                                } catch (com.fasterxml.jackson.core.JsonProcessingException e) {
                                    log.error("JSON parsing error for Ollama stream line: {}", line, e);
                                } catch (Exception e) {
                                    log.error("处理Ollama流数据块失败: {}", line, e);
                                    Map<String, String> errorEventMap = new HashMap<>();
                                    errorEventMap.put("type", "error");
                                    errorEventMap.put("content", "处理AI响应流时发生错误: " + e.getMessage());
                                    try {
                                        emitter.send(SseEmitter.event().data(objectMapper.writeValueAsString(errorEventMap)));
                                    } catch (Exception sendEx) { /* ignore */ }
                                }
                            }
                        } catch (java.io.IOException e) {
                            log.error("读取Ollama流响应失败", e);
                            Map<String, String> errorEventMap = new HashMap<>();
                            errorEventMap.put("type", "error");
                            errorEventMap.put("content", "读取AI服务响应流失败: " + e.getMessage());
                            try {
                                emitter.send(SseEmitter.event().data(objectMapper.writeValueAsString(errorEventMap)));
                            } catch (Exception sendEx) { /* ignore */ }
                            throw e; 
                        }
                        return null; 
                    }
                );
            } catch (org.springframework.web.client.ResourceAccessException e) {
                log.error("Ollama API连接超时或访问错误", e);
                sendSseEvent.accept("thought", "连接Ollama服务失败，可能是网络问题或API配置无效");
                sendSseEvent.accept("error", "连接Ollama服务失败: " + e.getMessage());
                throw e; // 继续抛出，让外层处理
            }

            emitter.complete();

        } catch (Exception e) {
            log.error("Ollama流式API调用失败", e);
            if (emitter != null && !emitter.toString().contains("completed")) { 
                 try {
                    // 提供更具体的错误信息
                    String errorMessage = e.getMessage();
                    if (e instanceof java.net.ConnectException || 
                        (e.getCause() != null && e.getCause() instanceof java.net.ConnectException)) {
                        errorMessage = "连接Ollama服务超时。请检查您的网络和API配置。";
                    }
                    
                    Map<String, String> errorEventMap = new HashMap<>();
                    errorEventMap.put("type", "thought");
                    errorEventMap.put("content", "AI服务连接失败，请检查网络和API配置");
                    emitter.send(SseEmitter.event().data(objectMapper.writeValueAsString(errorEventMap)));
                    
                    errorEventMap = new HashMap<>();
                    errorEventMap.put("type", "error");
                    errorEventMap.put("content", errorMessage);
                    emitter.send(SseEmitter.event().data(objectMapper.writeValueAsString(errorEventMap)));
                 } catch (Exception ex) {  }
                emitter.completeWithError(e);
            }
            
            // 创建一个描述性更强的错误消息
            AiChatMessage errorResult = new AiChatMessage();
            errorResult.setChatId(chat.getId());
            errorResult.setRole("assistant");
            
            String errorContent;
            if (e instanceof java.net.ConnectException || 
                (e.getCause() != null && e.getCause() instanceof java.net.ConnectException)) {
                errorContent = "Ollama服务连接超时。请检查您的网络连接和API配置。";
            } else {
                errorContent = "Ollama流式处理失败: " + e.getMessage();
            }
            
            errorResult.setContent(errorContent);
            errorResult.setTokenCount(0);
            errorResult.setStatus(1); 
            return errorResult;
        }
        
        
        // 保存完整的AI消息
        AiChatMessage finalAiMessage = new AiChatMessage();
        finalAiMessage.setChatId(chat.getId());
        finalAiMessage.setRole("assistant");
        finalAiMessage.setContent(fullResponseContent.toString().trim()); // Use the full content with thoughts and tags
        finalAiMessage.setTokenCount(totalEvalCount[0] > 0 ? totalEvalCount[0] : estimateTokens(fullResponseContent.toString()));
        finalAiMessage.setStatus(1);

        // Mark the stream as complete
        emitter.complete();

        return finalAiMessage; // This AiChatMessage is for saving to DB.
    }
    
    @Override
    public int countTokens(String content) {
        if (content == null || content.isEmpty()) {
            return 0;
        }
        
        try {
            // 使用cl100k_base编码，因为它适用于大多数现代模型
            // 包括GPT-3.5-turbo和GPT-4系列
            Encoding encoding = registry.getEncoding(EncodingType.CL100K_BASE);
            
            // 编码文本并计算token数量
            List<Integer> tokens = encoding.encode(content);
            return tokens.size();
        } catch (Exception e) {
            log.error("Token计数失败", e);
            // 如果编码失败，使用简单的估计方法
            return estimateTokens(content);
        }
    }
    
    /**
     * 备用的token估算方法，当编码库失败时使用
     */
    private int estimateTokens(String content) {
        // 英文平均每个单词约1.3个token
        // 中文等其他语言平均每个字符约1个token
        // 这是一个粗略的估计，实际可能有所不同
        String[] words = content.split("\\s+");
        return (int) Math.ceil(words.length * 1.3);
    }
} 